package com.sun.dream.fish.secret.asymmetry;

import javax.crypto.KeyAgreement;
import java.security.*;
import java.security.spec.InvalidKeySpecException;
import java.util.Base64;

public class DiffieHellmanExample {
    public static void main(String[] args) throws NoSuchAlgorithmException, InvalidKeySpecException, InvalidKeyException {
        // 1. 生成发送方和接收方的密钥对
        KeyPair senderKeyPair = generateKeyPair();
        KeyPair receiverKeyPair = generateKeyPair();

        // 2. 发送方使用自己的私钥和接收方的公钥生成共享密钥
        byte[] senderSharedSecret = generateSharedSecret(senderKeyPair.getPrivate(), receiverKeyPair.getPublic());

        // 3. 接收方使用自己的私钥和发送方的公钥生成共享密钥
        byte[] receiverSharedSecret = generateSharedSecret(receiverKeyPair.getPrivate(), senderKeyPair.getPublic());

        // 4. 打印生成的共享密钥（确保发送方和接收方的共享密钥相同）
        String senderSharedSecretBase64 = encodeSharedSecret(senderSharedSecret);
        String receiverSharedSecretBase64 = encodeSharedSecret(receiverSharedSecret);

        System.out.println("发送方的共享密钥：" + senderSharedSecretBase64);
        System.out.println("接收方的共享密钥：" + receiverSharedSecretBase64);
    }

    private static KeyPair generateKeyPair() throws NoSuchAlgorithmException {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("DH");
        keyPairGenerator.initialize(2048);

        return keyPairGenerator.generateKeyPair();
    }

    private static byte[] generateSharedSecret(PrivateKey privateKey, PublicKey publicKey) throws NoSuchAlgorithmException, InvalidKeyException {
        KeyAgreement keyAgreement = KeyAgreement.getInstance("DH");
        keyAgreement.init(privateKey);
        keyAgreement.doPhase(publicKey, true);

        return keyAgreement.generateSecret();
    }

    private static String encodeSharedSecret(byte[] sharedSecret) {
        return Base64.getEncoder().encodeToString(sharedSecret);
    }
}